package com.passion.source;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.passion.utils.DateBaseUtil;
import org.apache.flume.Context;
import org.apache.flume.EventDrivenSource;
import org.apache.flume.channel.ChannelProcessor;
import org.apache.flume.conf.Configurable;
import org.apache.flume.source.AbstractSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.*;


/**
 * @program flume-component
 * @Description 采集数据库数据
 * @Author limeihua
 * @Date 2024/09/02 15:45
 * @Version 1.0
 */
public class JdbcSource extends AbstractSource implements Configurable, EventDrivenSource {
    private static final Logger logger = LoggerFactory.getLogger(JdbcSource.class);
    private String username;
    private String password;
    private String database;
    private String queryFields;
    private Connection connection;
    private String url;
    private String tableName;
    private String separator;
    private String taskHost;
    private String taskName;
    private String condition;
    private String hostUsername;
    private String hostPassword;

    @Override
    public void configure(Context context) {
        // mysql的url
        url = context.getString("url", "jdbc:mysql://localhost:3306/");
        // mysql用户名
        username = context.getString("user", "root");
        // mysql密码
        password = context.getString("password", "159357");
        // 采集表在的数据库
        database = context.getString("database", "test");
        // 采集表名
        tableName = context.getString("tableName");
        // 采集查询字段
        queryFields = context.getString("queryFields", "*");
        // 每隔数值之间的分隔符
        separator = context.getString("separator", "\t");
        // 任务名称
        taskName = context.getString("taskName");
        // 任务所在节点
        taskHost = context.getString("taskHost");
        // 条件查询
        condition = context.getString("condition");
        // 节点用户名
        hostUsername = context.getString("hostUsername");
        // 节点密码
        hostPassword = context.getString("hostPassword");
    }

    @Override
    public synchronized void start() {
        try {
            ChannelProcessor channelProcessor = getChannelProcessor();
            connection = DriverManager.getConnection(url + database, username, password);
            Statement CStatement = connection.createStatement();
            ResultSet resultSetCount = CStatement.executeQuery("select count(*) count_row from " + tableName + " " + condition);
            resultSetCount.next();
            int rate = 0;
            int countRow = resultSetCount.getInt("count_row");
            logger.info("查询总条数为：" + countRow);
            if (countRow <= 500) {
                String executeSql = "select " + queryFields + " from " + tableName + " " + condition;
                logger.info("执行sql为：" + executeSql);
                DateBaseUtil.paginationExecute(executeSql, url + database, username, password, channelProcessor, separator);
            } else {
                while (countRow > 0) {
                    DateBaseUtil.paginationExecute("select " + queryFields + " from " + tableName + " " + condition + " limit " + rate * 500 + ",500", url + database, username, password, channelProcessor, separator);
                    countRow = countRow - 500;
                    rate = rate + 1;
                }
            }
        } catch (SQLException e) {
            throw new RuntimeException("Error executing SQL query");
        }
        // 由于是离线采集数据任务结束时需要将任务杀死
        try {
            JSch jsch = new JSch();
            // 访问节点
            Session session = jsch.getSession(hostUsername, taskHost, 22);
            session.setPassword(hostPassword);
            session.setConfig("StrictHostKeyChecking", "no");
            session.connect();
            String command = "ps aux | grep '" + taskName + "' | grep -v grep | awk '{print $2}' | xargs kill -9";
            logger.info("正在执行杀死任务命令：" + command);
            ChannelExec channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(command);
            channel.connect();
            channel.disconnect();
            session.disconnect();
        } catch (JSchException e) {
            throw new RuntimeException("任务杀死失败！！！");
        }
    }

    @Override
    public synchronized void stop() {
        if (connection != null) {
            try {
                connection.close();
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }
}


